"""
# Reference

[]()

This model is based on the following implementations:
- https://github.com/calmisential/TensorFlow2.0_ResNet
- https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py

"""
from tensorflow.keras import Model, layers, activations


class AlexNet(Model):

    def __init__(self, num_classes=10):
        super(AlexNet, self).__init__()
        self.conv1 = layers.Conv2D(96, 11, 4, 'valid', activation=activations.relu)
        self.max_pool1 = layers.MaxPool2D(3, 2, 'valid')
        self.bn1 = layers.BatchNormalization()
        self.conv2 = layers.Conv2D(256, 5, 1, 'same', activation=activations.relu)
        self.max_pool2 = layers.MaxPool2D(3, 2, 'same')
        self.bn2 = layers.BatchNormalization()
        self.conv3 = layers.Conv2D(384, 3, 1, 'same', activation=activations.relu)
        self.conv4 = layers.Conv2D(384, 3, 1, 'same', activation=activations.relu)
        self.conv5 = layers.Conv2D(256, 3, 1, 'same', activation=activations.relu)
        self.max_pool5 = layers.MaxPool2D(3, 2, 'same')
        self.bn5 = layers.BatchNormalization()

        self.flatten = layers.Flatten()
        self.fc1 = layers.Dense(4096, activation=activations.relu)
        self.dropout1 = layers.Dropout(0.2)
        self.fc2 = layers.Dense(4096, activation=activations.relu)
        self.dropout2 = layers.Dropout(0.2)
        self.out = layers.Dense(num_classes, activation=activations.softmax)

    def call(self, inputs, training=None, mask=None):
        x = self.conv1(inputs)
        x = self.max_pool1(x)
        x = self.bn1(x)

        x = self.conv2(x)
        x = self.max_pool2(x)
        x = self.bn2(x)

        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x = self.max_pool5(x)
        x = self.bn5(x)

        x = self.flatten(x)

        x = self.fc1(x)
        x = self.dropout1(x)
        x = self.fc2(x)
        x = self.dropout2(x)

        x = self.out(x)
        return x
